"""
图像识别模块

此模块使用讯飞星火大模型 API 进行图像内容识别，
能够分析上传的图片并返回可能的物品标签列表。

主要功能：
- 连接讯飞星火 WebSocket API
- 上传图片并获取 AI 识别结果
- 解析识别结果，返回物品标签列表

依赖模块：
- websocket: WebSocket 客户端
- PIL: 图像处理
- base64: 数据编码
- hmac, hashlib: 签名生成
- ssl: SSL 连接支持

注意：需要配置正确的 appid, api_key, api_secret 参数
"""

import base64
import json
import time
import ssl
import websocket
from PIL import Image
from io import BytesIO
from datetime import datetime
from time import mktime
from urllib.parse import urlparse, urlencode
from wsgiref.handlers import format_date_time
import hmac
import hashlib
import _thread as thread

# 配置参数
appid = ""
api_secret = ""
api_key = ""
imageunderstanding_url = "wss://spark-api.cn-huabei-1.xf-yun.com/v2.1/image"

# 全局变量存储识别结果
ai_result = ""
is_complete = False

class Ws_Param(object):
    """
    WebSocket 参数类

    用于生成讯飞星火 API 的认证参数和请求 URL。
    处理 API 密钥签名和授权头部的生成。
    """
    def __init__(self, APPID, APIKey, APISecret, url):
        """
        初始化 WebSocket 参数

        参数：
        APPID (str): 应用 ID
        APIKey (str): API 密钥
        APISecret (str): API 密钥
        url (str): WebSocket URL
        """
        self.APPID = APPID
        self.APIKey = APIKey
        self.APISecret = APISecret
        self.host = urlparse(url).netloc
        self.path = urlparse(url).path
        self.url = url

    def create_url(self):
        """
        生成带认证参数的 WebSocket URL

        使用 HMAC-SHA256 算法生成签名，并构建包含授权信息的完整 URL。

        返回：
        str: 带认证参数的 WebSocket URL
        """
        now = datetime.now()
        date = format_date_time(mktime(now.timetuple()))
        signature_origin = f"host: {self.host}\ndate: {date}\nGET {self.path} HTTP/1.1"
        signature_sha = hmac.new(self.APISecret.encode('utf-8'), signature_origin.encode('utf-8'),
                                 digestmod=hashlib.sha256).digest()
        signature_sha_base64 = base64.b64encode(signature_sha).decode(encoding='utf-8')
        authorization_origin = f'api_key="{self.APIKey}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature_sha_base64}"'
        authorization = base64.b64encode(authorization_origin.encode('utf-8')).decode(encoding='utf-8')
        v = {
            "authorization": authorization,
            "date": date,
            "host": self.host
        }
        return self.url + '?' + urlencode(v)

def on_error(ws, error):
    """WebSocket 错误回调函数"""
    print(f"### 错误: {error}")

def on_close(ws, one, two):
    """WebSocket 关闭回调函数"""
    global is_complete
    is_complete = True

def on_open(ws):
    """WebSocket 连接建立回调函数"""
    thread.start_new_thread(run, (ws,))

def run(ws, *args):
    """发送识别请求到 WebSocket"""
    data = json.dumps(gen_params(appid=ws.appid, question=ws.question))
    ws.send(data)

def on_message(ws, message):
    """WebSocket 消息接收回调函数"""
    global ai_result, is_complete
    data = json.loads(message)
    code = data['header']['code']
    if code != 0:
        print(f'请求错误: {code}, {data}')
        ws.close()
    else:
        choices = data["payload"]["choices"]
        status = choices["status"]
        content = choices["text"][0]["content"]
        ai_result += content
        if status == 2:
            ws.close()
            is_complete = True

def gen_params(appid, question):
    """
    生成 API 请求参数

    构建发送给讯飞星火 API 的参数字典，包含应用 ID、模型参数和消息内容。

    参数：
    appid (str): 应用 ID
    question (list): 消息列表，包含图片和文本内容

    返回：
    dict: API 请求参数字典
    """
    return {
        "header": {"app_id": appid},
        "parameter": {
            "chat": {
                "domain": "imagev3",
                "temperature": 0.5,
                "top_k": 4,
                "max_tokens": 2028,
                "auditing": "default"
            }
        },
        "payload": {"message": {"text": question}}
    }

def recognize_image_tags(image_path):
    """
    识别图片内容并返回物品标签列表

    此函数接收图片路径，使用讯飞星火大模型 API 进行图像识别，
    解析 AI 返回的结果，提取可能的物品名称标签。

    处理流程：
    1. 读取并预处理图片（转换为 RGB 格式，压缩尺寸，编码为 base64）
    2. 构建识别请求消息
    3. 建立 WebSocket 连接并发送请求
    4. 等待并接收 AI 响应
    5. 解析响应结果，返回标签列表（最多 5 个）

    参数：
    image_path (str): 图片文件的路径

    返回：
    list: 识别出的物品标签列表，每个标签为字符串。如果识别失败，返回空列表。
    """
    global ai_result, is_complete
    ai_result = ""
    is_complete = False

    try:
        with Image.open(image_path) as img:
            if img.mode in ('RGBA', 'LA', 'P'):
                img = img.convert('RGB')
            max_size = (1024, 1024)
            img.thumbnail(max_size, Image.LANCZOS)
            buffered = BytesIO()
            img.save(buffered, format="JPEG", quality=90)
            img_data = buffered.getvalue()
            img_base64 = str(base64.b64encode(img_data), 'utf-8')

        question = [
            {"role": "user", "content": img_base64, "content_type": "image"},
            {"role": "user", "content": "请识别这个物品的名称,要求尽量覆盖可能的叫法,比如(苹果手机,手机,苹果,电话,iPhone),"
                                        "按照逗号分隔,按出现频率排序"
                                        "如果画面中识别出多个物品,首先排除不可能丢失的东西(比如,背景,环境,装饰等),只保留可能丢失的物品"
            }
        ]

        wsParam = Ws_Param(appid, api_key, api_secret, imageunderstanding_url)
        websocket.enableTrace(False)
        wsUrl = wsParam.create_url()
        ws = websocket.WebSocketApp(wsUrl, on_message=on_message, on_error=on_error, on_close=on_close, on_open=on_open)
        ws.appid = appid
        ws.question = question

        ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})
        start_time = time.time()
        while not is_complete and (time.time() - start_time) < 10:
            time.sleep(0.1)

        if ai_result:
            return [tag.strip() for tag in ai_result.split("，") if tag.strip()][:5]
        else:
            print("大模型未返回有效结果")
            return []

    except Exception as e:
        print(f"大模型调用失败: {str(e)}")
        return []
